import querier.weibo.utils as utils
from querier.esquerier import ElasticSearchQuerier
from utils.province_code import PROVINCE_CODE

PROVINCE_DECODE = {v: k for k, v in PROVINCE_CODE.items()}


class WeiboUserInfoAnalysisQuerier(ElasticSearchQuerier):
    def __init__(self, es, index, doc_type):
        super(WeiboUserInfoAnalysisQuerier, self).__init__(es, index, doc_type)

    def _build_query(self, args):
        filters = args.get('filters')
        field = args.get('field')
        if field:
            query = self._gen_query(filters, field)
        else:
            query = self._gen_query_user_info(filters)
        return query, {}, {'filters': filters, 'field': field}

    def _build_result(self, es_result, param):
        field = param.get('field')
        if field:
            agg = es_result['aggregations']
            data = extract_result(agg, field)
            return {
                "values": {
                    'user_counts': data['doc_counts'],
                }
            }
        else:
            res = extract_user_info(es_result)
            return {'user_info': res}

    @staticmethod
    def _gen_query(filters, field):
        filter_clause = []
        filter_clause = utils.add_filter_clause2(filter_clause, filters, 'user_id', 'must')
        query = {
            "query": {
                "bool": {
                    "filter": filter_clause
                }
            },
            "aggs": {
                "user_hist": {
                    "terms": {
                        'field': field,
                    },
                    "aggs": {
                        "fans_count": {"sum": {"field": "likes", "missing": 0}},
                    }
                }
            },

            "size": 0
        }

        return query

    @staticmethod
    def _gen_query_user_info(filters):
        filter_clause = []
        filter_clause = utils.add_filter_clause2(filter_clause, filters, 'user_id', 'must')
        query = {
            "query": {
                "bool": {
                    "filter": filter_clause
                }
            }
        }

        return query


def extract_result(agg, field):
    buckets = agg['user_hist']['buckets']

    doc_counts = dict()
    fans_count = dict()
    if field == 'province':
        for b in buckets:
            k = province_decode(b['key'])
            doc_counts[k] = b['doc_count']
            fans_count[k] = b['fans_count']['value']
    else:
        for b in buckets:
            k = b['key']
            doc_counts[k] = b['doc_count']
            fans_count[k] = b['fans_count']['value']

    sorted_keys = sorted(doc_counts.keys())
    return {
        'doc_counts': [{'key': k, 'value': doc_counts[k]} for k in sorted_keys],
        'fans_count': [{'key': k, 'value': fans_count[k]} for k in sorted_keys],
    }


def extract_user_info(data):
    hits = data['hits']
    user_hits = hits['hits']
    kv = {}
    for b in user_hits:
        user_id = b['_source']['user_id']
        fans = b['_source']['fans']
        user_type = b['_source']['user_type']
        gender = b['_source']['gender']
        verified_type = b['_source']['verified_type']
        bi_followers = b['_source']['bi_followers']
        user_name = b['_source']['user_name']
        friends = b['_source']['friends']
        location = b['_source']['location']
        province = b['_source']['province']
        city = b['_source']['city']
        kv[user_id] = {'fans': fans, 'user_type': user_type, 'gender': gender,
                       'verified_type': verified_type, 'bi_followers': bi_followers,
                       'user_name': user_name, 'friends': friends,
                       'location': location, 'province': province, 'city': city}
    return kv


def province_decode(province_code):
    return PROVINCE_DECODE.get(province_code, '未知')
